#!/bin/bash

# 脚本：一键安装 Zsh, Oh My Zsh, 及常用插件 (重构版)

# --- 安全设置 ---
# -e: 如果任何命令返回非零退出状态，则立即退出脚本。
# -u: 如果尝试使用未设置的变量，则视为错误并立即退出。
# -o pipefail: 如果管道中的任何命令失败，则整个管道的退出状态为失败。
set -euo pipefail

# --- 颜色定义 ---
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# --- 插件定义 ---
# 格式: "plugin_name|gitee_repo_url|github_repo_url"
# 您可以在这里添加或修改插件
declare -a PLUGINS_TO_INSTALL=(
    "zsh-syntax-highlighting|https://gitee.com/yz0812/zsh-syntax-highlighting.git|https://github.com/zsh-users/zsh-syntax-highlighting.git"
    "zsh-autosuggestions|https://gitee.com/yz0812/zsh-autosuggestions.git|https://github.com/zsh-users/zsh-autosuggestions.git"
)

# --- 全局变量 ---
ZSHRC_FILE="$HOME/.zshrc"
OH_MY_ZSH_DIR="$HOME/.oh-my-zsh"
ZSH_CUSTOM_PLUGINS_DIR="${ZSH_CUSTOM:-$OH_MY_ZSH_DIR/custom}/plugins"
USER_MIRROR_PREFERENCE="github" # 默认为 GitHub, 会被用户选择覆盖

# --- 日志函数 ---
log_info() { echo -e "${GREEN}[INFO] $1${NC}"; }
log_warn() { echo -e "${YELLOW}[WARN] $1${NC}"; }
log_error() { echo -e "${RED}[ERROR] $1${NC}"; }
log_step() { echo -e "\n${YELLOW}--- $1 ---${NC}"; }

# --- 辅助函数 ---

# 检查命令是否存在
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# 检查并请求 sudo 权限 (如果需要)
ensure_sudo() {
    if [[ $EUID -ne 0 ]]; then
        if ! command_exists sudo; then
            log_error "sudo 命令未找到。此脚本的某些部分可能需要sudo权限。"
            log_error "请安装 sudo 或以 root 用户身份运行。"
            # exit 1 # 可以取消注释以强制退出，或者让脚本继续尝试非 sudo 操作
        fi
        log_info "某些操作可能需要sudo权限。"
        # sudo -v # 更新 sudo 时间戳，如果需要密码会提示
    fi
}

# 询问用户镜像偏好
ask_mirror_preference() {
    log_step "选择镜像源"
    echo -e "${YELLOW}您希望优先使用哪个镜像源来下载 Oh My Zsh 和插件？${NC}"
    echo "1) Gitee (中国大陆用户推荐)"
    echo "2) GitHub (国际用户推荐)"
    local choice
    while true; do
        read -r -p "请输入选项 (1 或 2，默认为 2 GitHub): " choice
        choice=${choice:-2} # Default to 2 if input is empty
        case "$choice" in
            1)
                USER_MIRROR_PREFERENCE="gitee"
                log_info "已选择 Gitee 作为优先镜像源。"
                break
                ;;
            2)
                USER_MIRROR_PREFERENCE="github"
                log_info "已选择 GitHub 作为优先镜像源。"
                break
                ;;
            *)
                log_error "无效输入，请输入 1 或 2。"
                ;;
        esac
    done
}


# --- 核心功能函数 ---

# 步骤 1: 安装依赖
install_dependencies() {
    log_step "步骤 1: 检查并安装依赖 (git, curl, wget, zsh)"
    local required_packages_str="git curl wget zsh"
    local packages_to_install=()
    local package_manager=""
    local pkg

    if command_exists yum; then
        package_manager="yum"
    elif command_exists apt-get; then
        package_manager="apt-get"
    else
        log_error "未知的包管理器。请手动安装依赖: $required_packages_str"
        return 1
    fi
    log_info "检测到包管理器: $package_manager"

    log_info "正在检查依赖项状态..."
    for pkg in $required_packages_str; do
        local is_installed=false
        if [[ "$package_manager" == "yum" ]]; then
            if rpm -q "$pkg" >/dev/null 2>&1; then
                is_installed=true
            fi
        elif [[ "$package_manager" == "apt-get" ]]; then
            if dpkg -s "$pkg" >/dev/null 2>&1; then
                is_installed=true
            fi
        fi

        if $is_installed; then
            log_info "依赖 '$pkg' 已安装。"
        else
            log_info "依赖 '$pkg' 未安装，将添加到安装列表。"
            packages_to_install+=("$pkg")
        fi
    done

    if [[ ${#packages_to_install[@]} -eq 0 ]]; then
        log_info "所有必需的依赖项 ($required_packages_str) 均已安装。"
        return 0
    else
        local packages_to_install_str="${packages_to_install[*]}"
        log_info "使用 $package_manager 安装以下缺失的依赖: $packages_to_install_str"

        local install_success=false
        if [[ "$package_manager" == "yum" ]]; then
            if  sudo yum install -y $packages_to_install_str; then
                install_success=true
            fi
        elif [[ "$package_manager" == "apt-get" ]]; then
            if sudo apt-get install -y $packages_to_install_str; then
                install_success=true
            fi
        fi

        if $install_success; then
            log_info "成功安装缺失的依赖: $packages_to_install_str"
        else
            log_error "依赖安装失败: $packages_to_install_str"
            return 1
        fi
    fi
    return 0
}

# 步骤 2: 安装 Oh My Zsh
install_oh_my_zsh() {
    log_step "步骤 2: 安装 Oh My Zsh"
    if [ -d "$OH_MY_ZSH_DIR" ]; then
        log_warn "Oh My Zsh 似乎已经安装在 $OH_MY_ZSH_DIR。跳过安装。"
        return 0
    fi

    local install_script_url_gitee="https://gitee.com/yz0812/ohmyzsh/raw/master/tools/install.sh"
    local install_script_url_github="https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh"
    local primary_install_url=""
    local secondary_install_url=""
    local chosen_install_url=""
    local fetch_cmd=""

    if [[ "$USER_MIRROR_PREFERENCE" == "gitee" ]]; then
        primary_install_url="$install_script_url_gitee"
        secondary_install_url="$install_script_url_github"
    else # Default to GitHub
        primary_install_url="$install_script_url_github"
        secondary_install_url="$install_script_url_gitee"
    fi

    log_info "将尝试从首选源 ($USER_MIRROR_PREFERENCE) 下载 Oh My Zsh 安装脚本..."
    if command_exists curl; then
        if curl --output /dev/null --silent --head --fail "$primary_install_url"; then
            chosen_install_url="$primary_install_url"
            fetch_cmd="curl -fsSL"
        elif curl --output /dev/null --silent --head --fail "$secondary_install_url"; then
            log_warn "首选源下载失败，尝试备用源..."
            chosen_install_url="$secondary_install_url"
            fetch_cmd="curl -fsSL"
        fi
    elif command_exists wget; then
        if wget --spider --quiet "$primary_install_url"; then
            chosen_install_url="$primary_install_url"
            fetch_cmd="wget -O-"
        elif wget --spider --quiet "$secondary_install_url"; then
            log_warn "首选源下载失败，尝试备用源..."
            chosen_install_url="$secondary_install_url"
            fetch_cmd="wget -O-"
        fi
    fi

    if [ -z "$chosen_install_url" ]; then
        log_error "curl 或 wget 未找到，或者无法访问 Oh My Zsh 安装脚本 (已尝试两个源)。"
        log_error "请检查网络连接或手动安装 Oh My Zsh。"
        return 1
    fi

    log_info "从 $chosen_install_url 下载并执行 Oh My Zsh 安装脚本..."
    if CHSH=yes RUNZSH=no KEEP_ZSHRC=yes sh -c "$($fetch_cmd $chosen_install_url)" "" --unattended; then
        log_info "Oh My Zsh 安装成功。"
    else
        log_error "Oh My Zsh 安装失败。"
        return 1
    fi
    return 0
}

# 克隆插件的辅助函数
clone_plugin() {
    local plugin_name="$1"
    local plugin_target_dir="$2"
    local gitee_repo_url="$3"
    local github_repo_url="$4"
    local primary_repo_url=""
    local secondary_repo_url=""
    local primary_source_name=""
    local secondary_source_name=""

    if [ -d "$plugin_target_dir" ]; then
        log_warn "插件 '$plugin_name' 似乎已存在于 $plugin_target_dir。跳过克隆。"
        return 0
    fi

    if [[ "$USER_MIRROR_PREFERENCE" == "gitee" ]]; then
        primary_repo_url="$gitee_repo_url"
        primary_source_name="Gitee"
        secondary_repo_url="$github_repo_url"
        secondary_source_name="GitHub"
    else # Default to GitHub
        primary_repo_url="$github_repo_url"
        primary_source_name="GitHub"
        secondary_repo_url="$gitee_repo_url"
        secondary_source_name="Gitee"
    fi

    mkdir -p "$(dirname "$plugin_target_dir")"

    log_info "正在尝试从 $primary_source_name 克隆插件 '$plugin_name'..."
    if git clone --depth 1 "$primary_repo_url" "$plugin_target_dir"; then
        log_info "插件 '$plugin_name' 从 $primary_source_name 安装成功。"
        return 0
    else
        log_warn "从 $primary_source_name 克隆 '$plugin_name' 失败。正在尝试 $secondary_source_name..."
        if git clone --depth 1 "$secondary_repo_url" "$plugin_target_dir"; then
            log_info "插件 '$plugin_name' 从 $secondary_source_name 安装成功。"
            return 0
        else
            log_error "无法从 $primary_source_name 或 $secondary_source_name 克隆插件 '$plugin_name'。"
            return 1
        fi
    fi
}

# 步骤 3: 安装插件
install_all_plugins() {
    log_step "步骤 3: 安装 Zsh 插件"
    if [ ${#PLUGINS_TO_INSTALL[@]} -eq 0 ]; then
        log_warn "没有定义要安装的插件。"
        return 0
    fi

    local all_plugins_installed_successfully=true
    for plugin_info in "${PLUGINS_TO_INSTALL[@]}"; do
        IFS='|' read -r name gitee_url github_url <<< "$plugin_info"
        local plugin_dir="$ZSH_CUSTOM_PLUGINS_DIR/$name"
        log_info "安装插件: $name..."
        if ! clone_plugin "$name" "$plugin_dir" "$gitee_url" "$github_url"; then
            all_plugins_installed_successfully=false
            log_warn "插件 '$name' 安装失败，但脚本将继续。"
        fi
    done

    if $all_plugins_installed_successfully; then
        log_info "所有已定义的插件都已尝试安装。"
    else
        log_warn "部分插件未能成功安装。请检查上面的错误信息。"
    fi
    return 0
}

# 步骤 4: 配置 ~/.zshrc
configure_zshrc() {
    log_step "步骤 4: 配置 $ZSHRC_FILE"

    if [ ! -f "$ZSHRC_FILE" ] && [ -f "$HOME/.zshrc.pre-oh-my-zsh" ]; then
        log_info "$ZSHRC_FILE 未找到，但找到了 $HOME/.zshrc.pre-oh-my-zsh。将其复制为 $ZSHRC_FILE。"
        cp "$HOME/.zshrc.pre-oh-my-zsh" "$ZSHRC_FILE"
    elif [ ! -f "$ZSHRC_FILE" ]; then
        log_warn "$ZSHRC_FILE 未找到。Oh My Zsh 可能未正确初始化或您选择不创建它。"
        log_info "将创建一个基本的 $ZSHRC_FILE..."
        echo "export ZSH=\"$OH_MY_ZSH_DIR\"" > "$ZSHRC_FILE"
        echo "ZSH_THEME=\"robbyrussell\"" >> "$ZSHRC_FILE"
        # 初始化一个空的 plugins 数组，如果不存在的话
        echo "plugins=()" >> "$ZSHRC_FILE"
        echo "source \$ZSH/oh-my-zsh.sh" >> "$ZSHRC_FILE"
        log_info "已创建基本的 $ZSHRC_FILE。"
    fi

    # 确保 plugins=(...) 行存在，如果不存在，添加默认的 (git)
    if ! grep -q -E "^\s*plugins=\([^)]*\)" "$ZSHRC_FILE"; then
        log_info "在 $ZSHRC_FILE 中未找到 'plugins=(...)' 行。正在添加 'plugins=(git)'..."
        # 尝试添加到 ZSH_THEME 之后，或者文件末尾
        if grep -q "ZSH_THEME=" "$ZSHRC_FILE"; then
            sed -i "/ZSH_THEME=/a plugins=(git)" "$ZSHRC_FILE"
        else
            echo -e "\nplugins=(git)" >> "$ZSHRC_FILE"
        fi
    fi

    # 添加插件到 plugins 数组
    local current_plugins
    current_plugins=$(grep -E "^\s*plugins=\(([^)]*)\)" "$ZSHRC_FILE" | sed -E "s/^\s*plugins=\(([^)]*)\)/\1/")
    local plugins_to_add_to_zshrc=()

    # 只有当当前插件中不存在时，才添加这些默认插件
    if ! echo "$current_plugins" | grep -q -w "git"; then
        plugins_to_add_to_zshrc+=("git")
    fi
    
    if ! echo "$current_plugins" | grep -q -w "history"; then
        plugins_to_add_to_zshrc+=("history")
    fi
    
    if ! echo "$current_plugins" | grep -q -w "z"; then
        plugins_to_add_to_zshrc+=("z")
    fi
    
    if ! echo "$current_plugins" | grep -q -w "docker"; then
        plugins_to_add_to_zshrc+=("docker")
    fi
    
    if ! echo "$current_plugins" | grep -q -w "docker-compose"; then
        plugins_to_add_to_zshrc+=("docker-compose")
    fi

    if ! echo "$current_plugins" | grep -q -w "history-substring-search"; then
        plugins_to_add_to_zshrc+=("history-substring-search")
    fi

    if ! echo "$current_plugins" | grep -q -w "zsh-syntax-highlighting"; then
        plugins_to_add_to_zshrc+=("zsh-syntax-highlighting")
    fi

    if ! echo "$current_plugins" | grep -q -w "zsh-autosuggestions"; then
        plugins_to_add_to_zshrc+=("zsh-autosuggestions")
    fi
    for plugin_info in "${PLUGINS_TO_INSTALL[@]}"; do
        IFS='|' read -r name <<< "$plugin_info"

        # 检查插件是否已在数组中
        if ! echo "$current_plugins" | grep -q -w "$name"; then
            # 检查插件目录是否存在，只有存在才添加
            if [ -d "$ZSH_CUSTOM_PLUGINS_DIR/$name" ]; then
                plugins_to_add_to_zshrc+=("$name")
                log_info "计划将插件 '$name' 添加到 $ZSHRC_FILE。"
            else
                log_warn "插件 '$name' 的目录 ($ZSH_CUSTOM_PLUGINS_DIR/$name) 未找到，不会添加到 $ZSHRC_FILE。"
            fi
        else
            log_info "插件 '$name' 似乎已在 $ZSHRC_FILE 中配置。"
        fi
    done

    # 更新 plugins 行
    # 使用 awk 来构建新的 plugins 行，确保唯一性和正确格式
    # 首先提取现有的所有插件
    local final_plugin_list_str
    final_plugin_list_str=$( (echo "$current_plugins"; printf "%s\n" "${plugins_to_add_to_zshrc[@]}") | awk '!seen[$0]++ {print $0}' | tr '\n' ' ' | sed 's/ $//' )

    log_info "最终的插件列表将是: ($final_plugin_list_str)"

    # 使用临时文件进行 sed 操作，以提高兼容性 (macOS sed vs GNU sed)
    local tmp_zshrc_file
    tmp_zshrc_file=$(mktemp)

    # 替换旧的 plugins 行
    awk -v plugins_str="$final_plugin_list_str" '
        BEGIN { found=0 }
        /^\s*plugins=\([^)]*\)/ {
            if (!found) {
                print "plugins=(" plugins_str ")"
                found=1
            }
            next
        }
        { print }
        END {
            if (!found) { # 如果文件中没有 plugins= 行 (理论上前面已处理，但作为保险)
                print "\nplugins=(" plugins_str ")"
            }
        }
    ' "$ZSHRC_FILE" > "$tmp_zshrc_file" && mv "$tmp_zshrc_file" "$ZSHRC_FILE"

    # 在文件末尾添加 autoload -Uz compinit && compinit -i（如果不存在）
    if ! grep -q "autoload -Uz compinit && compinit -i" "$ZSHRC_FILE"; then
        echo -e "\n# 初始化自动补全功能\nautoload -Uz compinit && compinit -i" >> "$ZSHRC_FILE"
        log_info "已添加自动补全功能初始化命令。"
    else
        log_info "自动补全功能初始化命令已存在，不再重复添加。"
    fi

    log_info "$ZSHRC_FILE 配置更新完成。请检查文件以确保格式正确。"

    return 0
}

# 步骤 5: 更改默认 Shell 为 Zsh
set_default_shell() {
    log_step "步骤 5: 检查并设置 Zsh 为默认 Shell"
    local zsh_path
    if ! zsh_path=$(which zsh); then
        log_error "未找到 zsh 执行路径。请确保 zsh 已正确安装并位于 PATH 中。"
        return 1
    fi

    local current_shell
    current_shell=$(getent passwd "$USER" | cut -d: -f7)

    if [ "$current_shell" == "$zsh_path" ]; then
        log_info "Zsh ($zsh_path) 已经是您的默认 Shell。"
        return 0
    fi

    log_info "当前默认 Shell: $current_shell"
    log_info "Zsh 路径: $zsh_path"

    if ! command_exists chsh; then
        log_error "chsh 命令未找到。无法自动更改默认 Shell。"
        log_warn "请手动将 Zsh ($zsh_path) 设置为您的默认 Shell。例如: chsh -s $zsh_path"
        return 1
    fi

    log_info "正在尝试将默认 Shell 更改为 Zsh..."
    if sudo chsh -s "$zsh_path" "$USER"; then
        log_info "默认 Shell 已通过 sudo 尝试更改为 Zsh。"
        log_warn "您需要注销并重新登录才能使更改生效。"
        return 0
    else
        log_warn "使用 sudo chsh 更改默认 Shell 失败。"
        log_info "尝试不使用 sudo (可能需要当前用户密码)..."
        if chsh -s "$zsh_path"; then
            log_info "默认 Shell 已尝试更改为 Zsh (无需 sudo)。"
            log_warn "您需要注销并重新登录才能使更改生效。"
            return 0
        else
            log_error "更改默认 Shell 失败。"
            log_warn "您可能需要手动运行 'chsh -s $zsh_path' (可能需要密码) 或 'sudo chsh -s $zsh_path $USER'。"
            return 1
        fi
    fi
}

# --- 主函数 ---
main() {
    log_info "开始一键安装 Zsh, Oh My Zsh 及插件..."
    ensure_sudo
    ask_mirror_preference # 新增：询问用户镜像偏好

    install_dependencies || { log_error "依赖安装失败，脚本终止。"; exit 1; }
    install_oh_my_zsh || { log_warn "Oh My Zsh 安装失败或已跳过，继续执行插件安装和配置..."; }
    install_all_plugins
    configure_zshrc || { log_error "Zshrc 配置失败，脚本终止。"; exit 1; }
    set_default_shell

    log_info "\n--- 安装和配置基本完成! ---"
    log_warn "请注意:"
    log_warn "1. 您需要注销并重新登录，或者重启您的终端，才能使所有更改 (特别是默认 Shell 和 Zsh 配置) 完全生效。"
    log_warn "2. 如果 Zsh 提示符或插件未立即生效，请先尝试重新登录。"
    log_warn "   如果问题依旧，请检查 '$ZSHRC_FILE' 文件中的插件配置是否正确，例如："
    log_warn "   plugins=(git zsh-syntax-highlighting zsh-autosuggestions)"
    log_warn "3. 某些 Oh My Zsh 主题 (如 agnoster) 可能需要安装 Powerline 字体才能正确显示。"
    log_warn "4. 如果遇到 'compaudit: insecure directories' 错误，可以尝试运行:"
    log_warn "   chmod -R 755 \"$ZSH_CUSTOM_PLUGINS_DIR\" && compinit -u"
    log_warn "   或者针对 Oh My Zsh 目录: chmod -R 755 \"$OH_MY_ZSH_DIR\" && compinit -u"
    log_info "-------------------------您需要注销并重新登录，或者重启您的终端，才能使所有更改 (特别是默认 Shell 和 Zsh 配置) 完全生效。------------------------"
    exit 0
}

# --- 脚本入口 ---
main "$@"